import logging
from itertools import chain

from .tokenization_utils import Tokenizer

logger = logging.getLogger(__name__)


class DialTokenizer(Tokenizer):
    def __init__(self,
                 label_pad_id: int = -100,
                 max_history: int = 15,
                 **kwargs):
        super(DialTokenizer, self).__init__(**kwargs)
        # 对 token 进行检查
        if self.cls_token is None:
            raise ValueError(f" cls token 不能 为空，请检查")
        if self.cls_token not in self.vocab:
            raise ValueError(f" cls token 不在词典中，请检查")
        if self.sep_token is None:
            raise ValueError(f" sep token 不能 为空，请检查")
        if self.sep_token not in self.vocab:
            raise ValueError(f" sep token 不在词典中，请检查")

        if self.pad_token is None:
            raise ValueError(f" pad token 不能 为空，请检查")
        if self.pad_token not in self.vocab:
            raise ValueError(f" pad token 不在词典中，请检查")

        self.label_pad_id = label_pad_id
        self.max_history = max_history

        self.input_ids_key = 'input_ids'
        self.type_ids_key = 'type_ids'
        self.trg_input_ids_key = 'trg_input_ids'
        self.label_ids_key = 'label_ids'

    def encode(self,
               dialogue: [str],
               return_type_ids=False,
               return_label_ids=True,
               return_dict=False,
               **kwargs):
        if isinstance(dialogue, str):
            dialogue = [dialogue]
        dialogue_ids = [self.convert_tokens_to_ids(self.tokenize(utterance)) for utterance in dialogue]

        return self.encode_ids(dialogue_ids=dialogue_ids,
                               return_type_ids=return_type_ids,
                               return_label_ids=return_label_ids,
                               return_dict=return_dict,
                               **kwargs)

    def encode_ids(self,
                   dialogue_ids: [[int]],
                   return_type_ids=False,
                   return_label_ids=True,
                   return_dict=False,
                   **kwargs):
        raise NotImplementedError(f"子类实现 encode_ids 方法")

    def output_dict(self, input_ids=None, type_ids=None, trg_input_ids=None, label_ids=None):
        output_dict = {self.input_ids_key: input_ids,
                       self.type_ids_key: type_ids,
                       self.trg_input_ids_key: trg_input_ids,
                       self.label_ids_key: label_ids}
        return {k: v for k, v in output_dict.items() if v is not None}

    @classmethod
    def output_tuple(cls, input_ids=None, type_ids=None, trg_input_ids=None, label_ids=None):
        output_list = [item for item in [input_ids, type_ids, trg_input_ids, label_ids] if item is not None]
        return output_list[0] if len(output_list) == 1 else output_list

    @classmethod
    def check_dialogue_ids(cls, dialogue_ids: [[int]]):
        if not isinstance(dialogue_ids, list):
            raise ValueError("dialogue_ids 必须为 list 请检查")
        if isinstance(dialogue_ids[0], int):  # 转化为2维数组
            dialogue_ids = [dialogue_ids]

        if len(dialogue_ids) <= 0:
            raise ValueError('未检测到有效对话数据，请检查')
        elif len(dialogue_ids) == 1:
            logger.warning("对话数量不足2句，请注意")

        return dialogue_ids

    def _history_ids_from_dialogue_ids_(self,
                                        dialogue_ids: [[int]],
                                        cut_response=True,
                                        add_special=True):
        """ 从历史对话中提取 history ids

            cut_response: 剪切掉最后回复句，
                True 默认最后一句为回复句子，将剪切掉
                False 全部句子都作为历史对话
            add_special: 添加特殊符
        """
        dialogue_ids = self.check_dialogue_ids(dialogue_ids)

        dialogue_ids = dialogue_ids[-self.max_history*2:]

        if cut_response and len(dialogue_ids) > 1:
            history_ids = dialogue_ids[:-1]
        else:
            history_ids = dialogue_ids
        history_ids = self._sample_ids_cut_(dialogue_ids=history_ids)

        if add_special:
            history_ids = [sub_ids + [self.sep_token_id] for sub_ids in history_ids]
            history_ids[0] = [self.cls_token_id] + history_ids[0]
        return history_ids

    @classmethod
    def _type_ids_from_history_ids_(cls, history_ids: [[int]]):
        """ 从 history_ids 中获取 类型 ids"""
        type_ids = []
        for i, sub_ids in enumerate(history_ids):
            type_id = 0 if i % 2 == 0 else 1
            type_ids += [type_id] * len(sub_ids)
        return type_ids

    def _sample_ids_cut_(self, dialogue_ids):
        """ 将对话进行裁剪 """
        new_sample_ids = []
        total_len = 0
        for i, line in enumerate(dialogue_ids[::-1]):
            cur_len = len(line) + 1 + int(i == 0)
            if (total_len + cur_len) >= self.max_len - 2:
                cut_id = self.max_len - 2 - total_len - int(i == 0)
                cut_line_ = line[: cut_id]
                new_sample_ids.append(cut_line_)
                break
            else:
                new_sample_ids.append(line)
            total_len += cur_len
        new_sample_ids = new_sample_ids[::-1]
        return new_sample_ids


class Seq2SeqDialTokenizer(DialTokenizer):
    """ seq2seq 的编码方式
    dialogue: [str]   ->  [text1, text2, text3, text4, .....]   需返回label_ids时，默认最后一句为response
    dialogue_ids: [[int]] -> [text1_ids, texts_ids, text3_ids, ...]  需返回label_ids时，默认最后一句为response_ids

    history_ids: = [text1_ids, texts_ids, text3_ids, ...]
        若不需返回 label_ids，则等于 dialogue_ids
        若需要返回 label_ids，则等于 dialogue_ids[:-1]（当dialogue_ids长度为1 的时候，等于 dialogue_ids）
    response_ids: = dialogue_ids[-1]

    编码: 返回label_ids时，最后一句为label(response, response_ids)
        input_ids = [cls] + text1_ids + [sep] + text2_ids + [sep] ......
        type_ids = [0] + [0]*len(text1_ids) + [0] + [1]*len(text1_id2) + [1] ...
        trg_input: [cls] + response_ids
        label_ids: response_ids + [sep]
    """
    def __init__(self, label_pad_id: int = -100, max_history: int = 15, max_len=30, **kwargs):
        super(Seq2SeqDialTokenizer, self).__init__(label_pad_id=label_pad_id,
                                                   max_history=max_history,
                                                   max_len=max_len,
                                                   **kwargs)

    def encode_ids(self, dialogue_ids: [[int]],
                   return_type_ids=False,
                   return_label_ids=True,
                   return_dict=False,
                   **kwargs):
        dialogue_ids = self.check_dialogue_ids(dialogue_ids)
        history_ids = self._history_ids_from_dialogue_ids_(dialogue_ids=dialogue_ids,
                                                           cut_response=True,
                                                           add_special=True)
        input_ids = list(chain(*history_ids))

        type_ids = None
        if return_type_ids:
            type_ids = self._type_ids_from_history_ids_(history_ids=history_ids)

        trg_input = None
        label_ids = None
        if return_label_ids:
            response_ids = dialogue_ids[-1]
            trg_input = [self.cls_token_id] + response_ids
            label_ids = response_ids + [self.sep_token_id]

        if return_dict:
            return self.output_dict(input_ids=input_ids,
                                    type_ids=type_ids,
                                    trg_input_ids=trg_input,
                                    label_ids=label_ids)
        else:
            return self.output_tuple(input_ids=input_ids,
                                     type_ids=type_ids,
                                     trg_input_ids=trg_input,
                                     label_ids=label_ids)


class GPT2DialTokenizer(DialTokenizer):
    def __init__(self, label_mask_history: bool = False, **kwargs):
        super(GPT2DialTokenizer, self).__init__(**kwargs)
        self.label_mask_history = label_mask_history

    def encode_ids(self,
                   dialogue_ids: [[int]],
                   return_type_ids=True,
                   return_label_ids=True,
                   return_dict=False,
                   **kwargs):
        dialogue_ids = self.check_dialogue_ids(dialogue_ids)
        history_ids = self._history_ids_from_dialogue_ids_(dialogue_ids=dialogue_ids,
                                                           cut_response=False,
                                                           add_special=True)
        input_ids = list(chain(*history_ids))

        type_ids = None
        if return_type_ids:
            type_ids = self._type_ids_from_history_ids_(history_ids=history_ids)

        label_ids = None
        if return_label_ids:
            if (self.label_mask_history is False) or (len(history_ids) <= 1):
                label_ids = input_ids
            else:
                label_ids = list(chain(*[[self.label_pad_id] * len(sub_line) for sub_line in history_ids[:-1]]))
                label_ids += history_ids[-1]

        if return_dict:
            return self.output_dict(input_ids=input_ids, type_ids=type_ids, label_ids=label_ids)
        else:
            return self.output_tuple(input_ids=input_ids, type_ids=type_ids, label_ids=label_ids)

